namespace NBody.KD
{
    using System;
    using NBody;
    
    /// A k-D tree, based on the implementation described in
    /// Friedman, Bentley, and Finkel (1977) and Arya and Mount (1993).
    public class Tree
    {
        private IDParticle[] bucket;
        private int b;
        
        private Node root;
        public Node Root
        {
            get { return root; }
        }
        
        public Tree(NBodySystem system, int bucket_size)
        {            
            // copy system
            bucket = new IDParticle[system.NumParts];
            for (int i = 0; i < system.NumParts; i++)
            {
                bucket[i] = new IDParticle(system[i]);
                bucket[i].ID = i;
            }
            b = bucket_size;
            
            Console.Error.WriteLine("Building a KD-Tree for {0} particles with a bucket size of {1}", bucket.Length, b);
            
            // array for the bounds of the node --
            // bounds[0,0] = lower x, bounds[0,1] = lower y, bounds[0,2] = lower z
            // bounds[1,0] = upper x, bounds[1,1] = upper y, bounds[1,2] = upper z
            double[,] bounds = new double[2,3];
            for (int i = 0; i < bucket.Length; i++)
            {
                if (bucket[i].X < bounds[0,0]) bounds[0,0] = bucket[i].X;
                if (bucket[i].X > bounds[1,0]) bounds[1,0] = bucket[i].X;
                if (bucket[i].Y < bounds[0,1]) bounds[0,1] = bucket[i].Y;
                if (bucket[i].Y > bounds[1,1]) bounds[1,1] = bucket[i].Y;
                if (bucket[i].Z < bounds[0,2]) bounds[0,2] = bucket[i].Z;
                if (bucket[i].Z > bounds[1,2]) bounds[1,2] = bucket[i].Z;
            }
            
            Console.Error.WriteLine("Bounds for system:");
            Console.Error.WriteLine("  x1 = {0}, x2 = {1}", bounds[0,0], bounds[1,0]);
            Console.Error.WriteLine("  y1 = {0}, y2 = {1}", bounds[0,1], bounds[1,1]);
            Console.Error.WriteLine("  z1 = {0}, z2 = {1}", bounds[0,2], bounds[1,2]);
            
            DateTime t1 = DateTime.Now;
            root = BuildNodes(0, system.NumParts, bounds);
            DateTime t2 = DateTime.Now;
            TimeSpan ts = t2.Subtract(t1);
            Console.Error.WriteLine("Tree built in {0}", ts);
            
            Console.Error.WriteLine("... done");
        }                
        
        // Recursively build the nodes of the tree.  This works
        // by first finding the dimension under which the data has the
        // most spread, and then splitting the data about the median
        // in that dimension.  BuildNodes() is then called on each half.
        // Once the size of the data is small enough, a leaf node is formed.
        private Node BuildNodes(int start, int end, double[,] bounds)
        {
            int size = end - start;
            if (size <= b)
            {
                return new LeafNode(start, end, bucket, bounds);
            }
            else
            {
                int d = 0;
                int k = start + (size - 1) / 2;
                double spread;
                double maxspread = 0.0;
                for (int j = 0; j < 3; j++)
                {
                    spread = Spreadest(j, start, end);
                    if (spread > maxspread)
                    {
                        maxspread = spread;
                        d = j;
                    }
                }
                double p = Median(d, k, start, end);
                double[,] left = (double[,])bounds.Clone();
                double[,] right = (double[,])bounds.Clone();
                
                left[1,d] = p;
                right[0,d] = p;
                return new SplitNode(d, p, BuildNodes(start, k+1, left), 
                                     BuildNodes(k+1, end, right), bounds);
            }
        }
        
        // Find the dimension of which the data has the most spread
        private double Spreadest(int j, int start, int end)
        {
            double min = bucket[start].Position(j);
            double max = min;
            for (int i = start + 1; i < end; i++)
            {
                if (bucket[i].Position(j) < min)
                    min = bucket[i].Position(j);
                if (bucket[i].Position(j) > max)
                    max = bucket[i].Position(j);
            }
            return max - min;
        }

        // Rearrange the particle order such that all particles
        // with a d'th coordinate value less than the k'th particle's
        // are lower in index, and vice versa.  This function
        // permanently alters the NBody::System.
        private double Median(int d, int k, int start, int end)
        {
            int left = start;
            int right = end - 1;
            int i, j;
            double x;
            IDParticle w;

            while (left < right)
            {
                x = bucket[k].Position(d); i = left; j = right;
                while (i < j)
                {
                    while (bucket[i].Position(d) < x)
                        i++;
                    while (x < bucket[j].Position(d))
                        j--;
                    if (i <= j)
                    {
                        w = bucket[i];
                        bucket[i] = bucket[j];
                        bucket[j] = w;
                        i++;
                        j--;
                    }
                }
                if (j < k) left = i;
                if (k < i) right = j;
            }
            return bucket[k].Position(d);
        }
        
        /// Find the m nearest neighbours of the position given.
        /// The nearest neighbours and squared distances are stored in
        /// the arrays particles and dist2, respectively.
        public void FindNearest(double[] position, int m, out Particle[] particles, out double[] dist2)
        {
            dist2 = new double[m];
            particles = new Particle[m];
                    
            PriorityQueue pq = new PriorityQueue(m);
            for (int i = 0; i < m; i++)
                pq.Push(-1, Double.PositiveInfinity);
        
            double[] off = new double[] {0, 0, 0};
                
            SearchInfo info = new SearchInfo(pq, position, off);
        
            root.FindNearest(0.0, ref info);
        
            for (int i = m-1; i >=0; i--)
            {
                if (pq.TopQueue == -1)
                    throw new NBodyException("FindNearest failed for unknown reason");
                particles[i] = bucket[pq.TopQueue];                    
                dist2[i] = pq.TopPriority;
                pq.Pop();
            }
        }
        
        /// Find the nearest neighbours of the tt'th particle in the tree.
        /// Uses brute force method, for testing purposes only -- this will
        /// be very slow.
        public void FindNearestBruteForce(double[] position, int m, out Particle[] particles, out double[] dist2)
        {
            particles = new Particle[m];
            dist2 = new double[m];
        
            // automatically fill queue with infinities ...
            PriorityQueue pq = new PriorityQueue(m);
            for (int i = 0; i < m; i++)
                pq.Push(-1, Double.PositiveInfinity);
                
            for (int i = 0; i < bucket.Length; i++)
            {
                double d2 = DistanceSqd(position, bucket[i]);
                if (d2 != 0.0 && d2 < pq.TopPriority)
                {
                    pq.Pop();
                    pq.Push(i, d2);
                }
            }
                
            for (int i = m-1; i >=0; i--)
            {
                if (pq.TopQueue == -1)
                    throw new ApplicationException("FindNearest failed for unknown reason");
                particles[i] = bucket[pq.TopQueue];
                dist2[i] = pq.TopPriority;
                pq.Pop();
            }
        }
        
        // Metric distance (squared).
        private static double DistanceSqd(double[] pos, Particle p2)
        {        
            return (pos[0] - p2.X) * (pos[0] - p2.X) +
                   (pos[1] - p2.Y) * (pos[1] - p2.Y) +
                   (pos[2] - p2.Z) * (pos[2] - p2.Z);
        }
        
        //  Calculate the density of each particle in the tree by
        // smoothing over the nearest m particles.
        public double[] CalculateDensity(int m)
        {
            PriorityQueue pq = new PriorityQueue(m);
            double furthest = Double.PositiveInfinity;
            double[] off = new double[3];
            
            Console.Error.WriteLine("Calculating density by smoothing over {0} particles", m);
            
            double[] rho = new double[bucket.Length];
            for (int i = 0; i < bucket.Length; i++)
            {            
                for (int j = 0; j < m; j++)
                    pq.Push(-1, furthest);
            
                for (int j = 0; j < 3; j++)
                    off[j] = 0.0;
            
                SearchInfo info = new SearchInfo(pq, bucket[i].PositionArray, off);
                
                root.FindNearest(0.0, ref info);
            
                double new_furthest = 0.0;
            
                double hi = 0.5 * Math.Sqrt(pq.TopPriority);
                for (int j = 0; j < m; j++)
                {
                    if (pq.TopQueue == -1)
                        throw new NBodyException("Calculate Density failed");
                    
                    double rij = Math.Sqrt(pq.TopPriority);
                    double Wij = 0.5 * W(rij, hi);
                               
                    int id = pq.TopQueue;
                    rho[bucket[i].ID] += (double)(Wij * bucket[id].Mass);
                    rho[bucket[id].ID] += (double)(Wij * bucket[i].Mass);
                
                    if (i < bucket.Length - 1)
                    {
                        double d = DistanceSqd(bucket[i+1].PositionArray, bucket[pq.TopQueue]);
                        if (d > new_furthest)
                            new_furthest = d;
                    }
                    pq.Pop();
                }
                furthest = new_furthest + 1e-10;
                
                if (i % (bucket.Length / 80) == 0)
                    Console.Error.Write("#");
            }
            
            Console.Error.WriteLine("\n... done");
            return rho;
        }
        
        // The smoothing kernel
        private static double W(double r, double h)
        {
            double rh = r / h;
            if (rh < 1.0)
                return 1.0 / (Math.PI * h * h * h) * 
                       (1.0 - 0.75 * (2.0 - rh) * (rh * rh));
            else
                return 1.0 / (Math.PI * h * h * h) * 0.25 * (2.0 - rh) * 
                       (2.0 - rh) * (2.0 - rh);
        }
        
        // Calculate the potential of each particle
        public double[] CalculatePotential(double theta, double eps)
        {
            double[] phi = new double[bucket.Length];
            
            DateTime t1 = DateTime.Now;
            
            for (int i = 0; i < bucket.Length; i++)
            {
                root.CalcPotential(bucket[i], theta, eps, ref phi[bucket[i].ID]);
            }

            DateTime t2 = DateTime.Now;
            TimeSpan ts = t2.Subtract(t1);
            Console.Error.WriteLine("Potential calculated in {0}", ts);
                
            return phi;
        } 
    }
}
